package com.tlgen.rpc.utils;

import com.alibaba.nacos.api.exception.NacosException;
import com.alibaba.nacos.api.naming.NamingFactory;
import com.alibaba.nacos.api.naming.NamingService;
import com.alibaba.nacos.api.naming.pojo.Instance;
import com.alibaba.nacos.common.utils.StringUtils;
import org.springframework.core.env.Environment;

import javax.servlet.http.HttpServletRequest;
import java.util.*;
import java.util.function.Function;

public class LoadBalanceUtils {

    private static final HashMap<String, Function<HashMap<String, Integer>, String>> lbMap = new HashMap<>();

    static {
        lbMap.put("com.netflix.loadbalancer.RoundRobinRule", LoadBalanceUtils::roundRobin);
        lbMap.put("com.netflix.loadbalancer.RandomRule", LoadBalanceUtils::random);
        lbMap.put("com.netflix.loadbalancer.HashRule", LoadBalanceUtils::hash);
        lbMap.put("com.netflix.loadbalancer.WeightRandomRule", LoadBalanceUtils::weightRandom);
        lbMap.put("com.netflix.loadbalancer.weightRoundRobinRule", LoadBalanceUtils::weightRoundRobin);
//        loadBalanceMap.put("com.netflix.loadbalancer.WeightedResponseTimeRule", LoadBalanceUtils::random);
//        loadBalanceMap.put("com.netflix.loadbalancer.BestAvailableRule", LoadBalanceUtils::random);
//        loadBalanceMap.put("com.netflix.loadbalancer.AvailabilityFilteringRule", LoadBalanceUtils::random);
//        loadBalanceMap.put("com.netflix.loadbalancer.ZoneAvoidanceRule", LoadBalanceUtils::random);
    }

    public static String loadBalanceServer(String serviceName, String url) {
        HashMap<String, Integer> instanceMap = new HashMap<>();
        String serverPort;
        if (StringUtils.isNotBlank(url))
            return url;
        try {
            // 从 Nacos 拉取服务信息
            Environment environment = SpringUtils.getBean(Environment.class);
            selectNacosInstances(serviceName, environment, instanceMap);
            // 获取用户自定义的负载均衡策略配置
            String loadBalanceType = environment.getProperty(serviceName + ".ribbon.NFLoadBalancerRuleClassName");
            Function<HashMap<String, Integer>, String> function = lbMap.get(loadBalanceType);
            if (Objects.nonNull(function))
                serverPort = lbMap.get(loadBalanceType).apply(instanceMap);
            else
                serverPort = random(instanceMap);
            } catch(NacosException e){
                throw new RuntimeException(e);
            }
            return serverPort;
    }

    private static void selectNacosInstances(String serviceName, Environment environment,
                                             HashMap<String, Integer> serverWeightMap) throws NacosException {
        String serverAddr = environment.getProperty("spring.cloud.nacos.discovery.server-addr");
        String namespace = environment.getProperty("spring.cloud.nacos.discovery.namespace");
        String username = environment.getProperty("spring.cloud.nacos.discovery.username");
        String password = environment.getProperty("spring.cloud.nacos.discovery.password");
        Properties properties = new Properties();
        properties.put("serverAddr", serverAddr);
        properties.put("namespace", namespace);
        properties.put("username", username);
        properties.put("password", password);
        NamingService namingService = NamingFactory.createNamingService(properties);
        List<Instance> instances = namingService.getAllInstances(serviceName);
        for (Instance instance : instances) {
            serverWeightMap.put(instance.getIp() + ":" + instance.getPort(), 1);
        }
    }

    private static String getRemoteIp() {
        HttpServletRequest request = HttpUtils.getRequest();
        return request.getRemoteHost() + request.getRemotePort();
    }

    private static String random(HashMap<String, Integer> map) {
        // 重建一个 Map，避免服务器的上下线导致的并发问题
        Map<String, Integer> serverMap =
                new HashMap<>(map);

        // 取得 Ip 地址 List
        Set<String> keySet = serverMap.keySet();
        ArrayList<String> keyList = new ArrayList<String>(keySet);

        java.util.Random random = new java.util.Random();
        int randomPos = random.nextInt(keyList.size());

        return keyList.get(randomPos);
    }

    private static String hash(HashMap<String, Integer> serverMap) {
        Set<String> keySet = serverMap.keySet();
        ArrayList<String> keyList = new ArrayList<String>(keySet);

        int hashCode = getRemoteIp().hashCode();
        int serverListSize = keyList.size();
        int serverPos = hashCode % serverListSize;

        return keyList.get(serverPos);
    }

    private static String roundRobin(HashMap<String, Integer> serverMap) {
        Integer pos = 0;

        // 取得 Ip 地址 List
        Set<String> keySet = serverMap.keySet();
        ArrayList<String> keyList = new ArrayList<String>(keySet);

        String server = null;
        synchronized (pos) {
            server = keyList.get(pos);
            pos++;
        }
        return server;
    }

    private static String weightRandom(HashMap<String, Integer> serverMap) {
        // 取得 Ip 地址 List
        Set<String> keySet = serverMap.keySet();
        Iterator<String> iterator = keySet.iterator();

        List<String> serverList = new ArrayList<String>();
        while (iterator.hasNext()) {
            String server = iterator.next();
            int weight = serverMap.get(server);
            for (int i = 0; i < weight; i++)
                serverList.add(server);
        }

        java.util.Random random = new java.util.Random();
        int randomPos = random.nextInt(serverList.size());

        return serverList.get(randomPos);
    }

    private static String weightRoundRobin(HashMap<String, Integer> serverMap) {
        Integer pos = 0;

        // 取得 Ip 地址 List
        Set<String> keySet = serverMap.keySet();
        Iterator<String> iterator = keySet.iterator();

        List<String> serverList = new ArrayList<String>();
        while (iterator.hasNext()) {
            String server = iterator.next();
            int weight = serverMap.get(server);
            for (int i = 0; i < weight; i++)
                serverList.add(server);
        }

        String server = null;
        synchronized (pos) {
            server = serverList.get(pos);
            pos++;
        }

        return server;
    }

}
